{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🔥 LightlyTrain - Object Detection with DINOv3 LTDETR 🔥\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for object detection with our state-of-the-art LTDETR model built on [DINOv3](https://github.com/facebookresearch/dinov3).\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/object_detection.ipynb)\n",
    "\n",
    "> **Important**: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to `Runtime` > `Change runtime type` and selecting a GPU hardware accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "LightlyTrain can be installed directly via `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Important**: LightlyTrain is officially supported on\n",
    "> - Linux: CPU or CUDA\n",
    "> - MacOS: CPU only\n",
    "> - Windows (experimental): CPU or CUDA\n",
    ">\n",
    "> We are planning to support MPS for MacOS.\n",
    ">\n",
    "> Check the [installation instructions](https://docs.lightly.ai/train/stable/installation.html) for more details on installation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction using LightlyTrain's model weights\n",
    "\n",
    "### Download an example image\n",
    "\n",
    "Download an example image for inference with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000577932.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the model weights\n",
    "\n",
    "Then load the model weights with LightlyTrain's `load_model` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "model = lightly_train.load_model(\"dinov3/convnext-tiny-ltdetr-coco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict the objects\n",
    "\n",
    "And finally run `model.predict` on the image. The method accepts file paths, URLs,\n",
    "PIL Images, or tensors as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = model.predict(\"image.jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the results\n",
    "\n",
    "Finally, we visualize the image and results to check what objects were detected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.io import read_image\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "\n",
    "image = read_image(\"image.jpg\")\n",
    "image_with_boxes = draw_bounding_boxes(\n",
    "    image,\n",
    "    boxes=prediction[\"bboxes\"],\n",
    "    labels=[model.classes[label.item()] for label in prediction[\"labels\"]],\n",
    ")\n",
    "plt.imshow(image_with_boxes.permute(1, 2, 0))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The predicted boxes are in the absolute (x_min, y_min, x_max, y_max) format, i.e. represent coordinates of the bounding boxes in pixels.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train object detection model\n",
    "\n",
    "Training your own detection model is straightforward with LightlyTrain.\n",
    "\n",
    "### Download dataset\n",
    "\n",
    "First download a dataset, it must be in YOLO format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O coco128.zip https://www.kaggle.com/api/v1/datasets/download/ultralytics/coco128 && unzip coco128.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then start the training with the `train_object_detection` function. You can specify\n",
    "various training parameters such as the model architecture, number of training steps,\n",
    "batch size, learning rate, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_object_detection(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"dinov3/convnext-tiny-ltdetr-coco\",\n",
    "    steps=100,  # Small number of steps for demonstration, default is 90_000.\n",
    "    batch_size=4,  # Small batch size for demonstration, default is 16.\n",
    "    data={\n",
    "        \"path\": \"coco128\",\n",
    "        \"train\": \"images/train2017\",\n",
    "        \"val\": \"images/val2017\",\n",
    "        \"names\": {\n",
    "            0: \"person\",\n",
    "            1: \"bicycle\",\n",
    "            2: \"car\",\n",
    "            3: \"motorcycle\",\n",
    "            4: \"airplane\",\n",
    "            5: \"bus\",\n",
    "            6: \"train\",\n",
    "            7: \"truck\",\n",
    "            8: \"boat\",\n",
    "            9: \"traffic light\",\n",
    "            10: \"fire hydrant\",\n",
    "            11: \"stop sign\",\n",
    "            12: \"parking meter\",\n",
    "            13: \"bench\",\n",
    "            14: \"bird\",\n",
    "            15: \"cat\",\n",
    "            16: \"dog\",\n",
    "            17: \"horse\",\n",
    "            18: \"sheep\",\n",
    "            19: \"cow\",\n",
    "            20: \"elephant\",\n",
    "            21: \"bear\",\n",
    "            22: \"zebra\",\n",
    "            23: \"giraffe\",\n",
    "            24: \"backpack\",\n",
    "            25: \"umbrella\",\n",
    "            26: \"handbag\",\n",
    "            27: \"tie\",\n",
    "            28: \"suitcase\",\n",
    "            29: \"frisbee\",\n",
    "            30: \"skis\",\n",
    "            31: \"snowboard\",\n",
    "            32: \"sports ball\",\n",
    "            33: \"kite\",\n",
    "            34: \"baseball bat\",\n",
    "            35: \"baseball glove\",\n",
    "            36: \"skateboard\",\n",
    "            37: \"surfboard\",\n",
    "            38: \"tennis racket\",\n",
    "            39: \"bottle\",\n",
    "            40: \"wine glass\",\n",
    "            41: \"cup\",\n",
    "            42: \"fork\",\n",
    "            43: \"knife\",\n",
    "            44: \"spoon\",\n",
    "            45: \"bowl\",\n",
    "            46: \"banana\",\n",
    "            47: \"apple\",\n",
    "            48: \"sandwich\",\n",
    "            49: \"orange\",\n",
    "            50: \"broccoli\",\n",
    "            51: \"carrot\",\n",
    "            52: \"hot dog\",\n",
    "            53: \"pizza\",\n",
    "            54: \"donut\",\n",
    "            55: \"cake\",\n",
    "            56: \"chair\",\n",
    "            57: \"couch\",\n",
    "            58: \"potted plant\",\n",
    "            59: \"bed\",\n",
    "            60: \"dining table\",\n",
    "            61: \"toilet\",\n",
    "            62: \"tv\",\n",
    "            63: \"laptop\",\n",
    "            64: \"mouse\",\n",
    "            65: \"remote\",\n",
    "            66: \"keyboard\",\n",
    "            67: \"cell phone\",\n",
    "            68: \"microwave\",\n",
    "            69: \"oven\",\n",
    "            70: \"toaster\",\n",
    "            71: \"sink\",\n",
    "            72: \"refrigerator\",\n",
    "            73: \"book\",\n",
    "            74: \"clock\",\n",
    "            75: \"vase\",\n",
    "            76: \"scissors\",\n",
    "            77: \"teddy bear\",\n",
    "            78: \"hair drier\",\n",
    "            79: \"toothbrush\",\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the training is completely, the final model checkpoint is saved in `out/my_experiment/exported_models/exported_last.pt`. You can load it from there to perform inference.\n",
    "\n",
    "If you have a validation dataset, the best model according to the validation metric is saved in `out/my_experiment/exported_models/exported_best.pt` instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lightly_train.load_model(\"out/my_experiment/exported_models/exported_last.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = model.predict(\"image.jpg\")\n",
    "\n",
    "image = read_image(\"image.jpg\")\n",
    "image_with_boxes = draw_bounding_boxes(\n",
    "    image,\n",
    "    boxes=prediction[\"bboxes\"],\n",
    "    labels=[model.classes[label.item()] for label in prediction[\"labels\"]],\n",
    ")\n",
    "plt.imshow(image_with_boxes.permute(1, 2, 0))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
